"""
FastAPI 内容安全检测服务
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any
import sys
import os

# 添加validators目录到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from guardrails import Guard, OnFailAction
from validators.ml_content_safety import MLContentSafety
from validators.rule_based_safety import RuleBasedSafety

# 创建FastAPI应用
app = FastAPI(
    title="内容安全检测API",
    description="基于BERT模型的LLM输出内容安全检测服务",
    version="1.0.0"
)

# 全局验证器实例(启动时初始化一次,避免重复加载模型)
global_ml_validator = None
global_rule_validator = None

@app.on_event("startup")
async def startup_event():
    """启动时加载模型"""
    global global_ml_validator, global_rule_validator

    print("🚀 初始化验证器...")

    # 初始化规则验证器
    global_rule_validator = RuleBasedSafety(on_fail=OnFailAction.EXCEPTION)
    print("✅ 规则验证器加载完成")

    # 初始化ML验证器
    global_ml_validator = MLContentSafety(
        model_path="../models/safety_classifier_v5_roberta",
        threshold=0.5,
        on_fail=OnFailAction.NOOP
    )
    print("✅ ML模型验证器加载完成")

# 请求模型
class ValidationRequest(BaseModel):
    """验证请求"""
    text: str = Field(..., description="待验证的文本内容")
    threshold: float = Field(0.5, ge=0.0, le=1.0, description="置信度阈值")
    model_path: Optional[str] = Field(
        "../models/safety_classifier",
        description="模型路径"
    )

# 响应模型
class ValidationResponse(BaseModel):
    """验证响应"""
    is_safe: bool = Field(..., description="内容是否安全")
    confidence: float = Field(..., description="模型置信度")
    predicted_class: int = Field(..., description="预测类别 (0:安全, 1:危险)")
    message: str = Field(..., description="验证消息")
    validated_output: Optional[str] = Field(None, description="验证后的输出")
    error: Optional[str] = Field(None, description="错误信息")

# 健康检查
@app.get("/")
def root():
    """根路径 - API信息"""
    return {
        "service": "内容安全检测API",
        "version": "1.0.0",
        "status": "运行中",
        "endpoints": {
            "验证": "/api/validate",
            "批量验证": "/api/validate/batch",
            "健康检查": "/health"
        }
    }

@app.get("/health")
def health_check():
    """健康检查"""
    return {"status": "healthy"}

@app.post("/api/validate", response_model=ValidationResponse)
def validate_content(request: ValidationRequest):
    """
    验证文本内容安全性

    Args:
        request: 验证请求

    Returns:
        ValidationResponse: 验证结果
    """
    if global_ml_validator is None or global_rule_validator is None:
        raise HTTPException(status_code=500, detail="验证器未初始化")

    try:
        # 创建Guardrails Guard - 使用全局验证器实例
        # 第一层: 基于规则的快速前置过滤
        # 第二层: 基于BERT模型的深度检测
        guard = Guard().use(global_rule_validator).use(global_ml_validator)

        # 执行验证
        result = guard.validate(request.text)

        # 构建响应
        if result.validation_passed:
            # 安全内容
            metadata = result.validated_output if hasattr(result, 'validated_output') else {}
            if isinstance(metadata, dict):
                predicted_class = metadata.get('predicted_class', 0)
                confidence = metadata.get('confidence', 1.0)
            else:
                predicted_class = 0
                confidence = 1.0

            return ValidationResponse(
                is_safe=True,
                confidence=confidence,
                predicted_class=predicted_class,
                message="内容安全，可以输出",
                validated_output=request.text,
                error=None
            )
        else:
            # 危险内容
            error_msg = result.error if hasattr(result, 'error') else "检测到违规内容"

            # 尝试从metadata获取详细信息
            predicted_class = 1
            confidence = 0.5

            if hasattr(result, 'metadata') and result.metadata:
                predicted_class = result.metadata.get('predicted_class', 1)
                confidence = result.metadata.get('confidence', 0.5)

            return ValidationResponse(
                is_safe=False,
                confidence=confidence,
                predicted_class=predicted_class,
                message="内容包含违规信息，不应输出",
                validated_output=None,
                error=error_msg
            )

    except Exception as e:
        import traceback

        # 检查是否是规则层ValidationError
        error_class_name = type(e).__name__

        if error_class_name == 'ValidationError':
            # 规则层拦截,返回危险响应而不是500错误
            error_msg = str(e)
            print(f"\n⛔ 规则层拦截: {error_msg}")

            return ValidationResponse(
                is_safe=False,
                confidence=1.0,
                predicted_class=1,
                message="内容包含违规信息（规则层拦截）",
                validated_output=None,
                error=error_msg
            )

        # 其他错误类型
        if isinstance(e, FileNotFoundError):
            error_detail = f"模型文件未找到: {request.model_path}"
            print(f"\n❌ FileNotFoundError: {error_detail}")
            print(traceback.format_exc())
            raise HTTPException(
                status_code=404,
                detail=error_detail
            )
        else:
            error_detail = f"验证失败: {str(e)}"
            print(f"\n❌ Exception in validate_content:")
            print(f"   错误类型: {type(e).__name__}")
            print(f"   错误信息: {str(e)}")
            print(f"   完整堆栈:")
            print(traceback.format_exc())
            raise HTTPException(
                status_code=500,
                detail=error_detail
            )

class BatchValidationRequest(BaseModel):
    """批量验证请求"""
    texts: list[str] = Field(..., description="待验证的文本列表")
    threshold: float = Field(0.5, ge=0.0, le=1.0, description="置信度阈值")
    model_path: Optional[str] = Field(
        "../models/safety_classifier",
        description="模型路径"
    )

class BatchValidationResponse(BaseModel):
    """批量验证响应"""
    results: list[ValidationResponse] = Field(..., description="验证结果列表")
    summary: Dict[str, Any] = Field(..., description="统计摘要")

@app.post("/api/validate/batch", response_model=BatchValidationResponse)
def validate_batch(request: BatchValidationRequest):
    """
    批量验证文本内容

    Args:
        request: 批量验证请求

    Returns:
        BatchValidationResponse: 批量验证结果
    """
    results = []
    safe_count = 0
    unsafe_count = 0

    for text in request.texts:
        single_request = ValidationRequest(
            text=text,
            threshold=request.threshold,
            model_path=request.model_path
        )

        result = validate_content(single_request)
        results.append(result)

        if result.is_safe:
            safe_count += 1
        else:
            unsafe_count += 1

    return BatchValidationResponse(
        results=results,
        summary={
            "total": len(request.texts),
            "safe": safe_count,
            "unsafe": unsafe_count,
            "safe_percentage": (safe_count / len(request.texts) * 100) if request.texts else 0
        }
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
